import time
import threading

# Snowflake ID 的各部分 bit 位数定义
WORKER_ID_BITS = 10  # 机器ID占10位（最多支持1024个节点）
SEQUENCE_BITS = 12  # 序列号占12位（每毫秒最多生成4096个ID）

# 最大值（通过位移计算）
MAX_WORKER_ID = -1 ^ (-1 << WORKER_ID_BITS)  # 1023
MAX_SEQUENCE = -1 ^ (-1 << SEQUENCE_BITS)  # 4095

# 左移位数（用于组成最终ID）
WORKER_ID_SHIFT = SEQUENCE_BITS  # 机器ID左移12位
TIMESTAMP_LEFT_SHIFT = SEQUENCE_BITS + WORKER_ID_BITS  # 时间戳左移22位

# 起始时间戳（2020-01-01），毫秒级
EPOCH = 1577836800000


class Snowflake:
    def __init__(self, worker_id: int):
        """
        初始化 Snowflake 生成器
        :param worker_id: 当前机器ID（必须保证集群中唯一）
        """
        if worker_id > MAX_WORKER_ID or worker_id < 0:
            raise ValueError(f"worker_id必须在0到{MAX_WORKER_ID}之间")

        self.worker_id = worker_id  # 当前节点ID
        self.sequence = 0  # 当前毫秒内的序列号
        self.last_timestamp = -1  # 上一次生成ID的时间戳
        self.lock = threading.Lock()  # 保证线程安全

    def _current_millis(self):
        """返回当前时间的毫秒数"""
        return int(time.time() * 1000)

    def _wait_until_next_millis(self, last_timestamp):
        """
        如果当前时间小于等于上一次时间，阻塞直到下一毫秒
        这是为了避免在同一毫秒内序列号溢出
        """
        timestamp = self._current_millis()
        while timestamp <= last_timestamp:
            timestamp = self._current_millis()
        return timestamp

    def get_id(self):
        """生成一个全局唯一ID"""
        with self.lock:  # 加锁防止并发冲突
            timestamp = self._current_millis()

            # 如果当前时间小于上次时间，说明系统时钟回拨，抛出异常
            if timestamp < self.last_timestamp:
                raise Exception(f"时钟回拨，拒绝生成ID，当前：{timestamp}，上次：{self.last_timestamp}")

            if timestamp == self.last_timestamp:
                # 同一毫秒内，序列号自增
                self.sequence = (self.sequence + 1) & MAX_SEQUENCE
                if self.sequence == 0:
                    # 当前毫秒内的序列号用完，等待下一毫秒
                    timestamp = self._wait_until_next_millis(self.last_timestamp)
            else:
                # 不同毫秒内，序列号重置
                self.sequence = 0

            self.last_timestamp = timestamp

            # 通过位移拼接出最终的64位ID
            id_ = ((timestamp - EPOCH) << TIMESTAMP_LEFT_SHIFT) | \
                  (self.worker_id << WORKER_ID_SHIFT) | \
                  self.sequence

            return id_


if __name__ == '__main__':
    # 创建一个 Snowflake 实例，worker_id=1
    snowflake = Snowflake(worker_id=1)

    # 连续生成几个 ID
    for _ in range(10):
        unique_id = snowflake.get_id()
        print(unique_id)
